package org.nhindirect.stagent;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;


import java.io.BufferedInputStream;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.InputStream;
import java.security.cert.CertStore;
import java.security.cert.CollectionCertStoreParameters;
import java.security.cert.X509Certificate;

import javax.mail.internet.ContentType;
import javax.mail.internet.InternetHeaders;
import javax.mail.internet.MimeBodyPart;
import javax.mail.internet.MimeMultipart;
import javax.mail.internet.ParseException;

import org.bouncycastle.asn1.ASN1EncodableVector;
import org.bouncycastle.asn1.DERObjectIdentifier;
import org.bouncycastle.asn1.cms.AttributeTable;
import org.bouncycastle.asn1.cryptopro.CryptoProObjectIdentifiers;
import org.bouncycastle.asn1.nist.NISTObjectIdentifiers;
import org.bouncycastle.asn1.oiw.OIWObjectIdentifiers;
import org.bouncycastle.asn1.pkcs.PKCSObjectIdentifiers;
import org.bouncycastle.asn1.smime.SMIMECapabilitiesAttribute;
import org.bouncycastle.asn1.smime.SMIMECapability;
import org.bouncycastle.asn1.smime.SMIMECapabilityVector;
import org.bouncycastle.asn1.teletrust.TeleTrusTObjectIdentifiers;
import org.bouncycastle.cms.CMSProcessable;
import org.bouncycastle.cms.CMSProcessableByteArray;
import org.bouncycastle.cms.CMSSignedData;
import org.bouncycastle.cms.CMSSignedDataGenerator;
import org.bouncycastle.cms.RecipientId;
import org.bouncycastle.cms.RecipientInformation;
import org.bouncycastle.cms.RecipientInformationStore;
import org.bouncycastle.cms.SignerInformation;
import org.bouncycastle.mail.smime.CMSProcessableBodyPartInbound;
import org.bouncycastle.mail.smime.SMIMEEnveloped;
import org.bouncycastle.mail.smime.SMIMEEnvelopedGenerator;
import org.bouncycastle.mail.smime.SMIMESignedGenerator;
import org.nhindirect.stagent.ProtocolException.ProtocolError;
import org.nhindirect.stagent.cert.X509CertificateEx;

import org.apache.commons.codec.binary.Base64;

public class Cryptographer 
{
	
    public static final String  DIGEST_SHA1 = OIWObjectIdentifiers.idSHA1.getId();
    public static final String  DIGEST_MD5 = PKCSObjectIdentifiers.md5.getId();
    public static final String  DIGEST_SHA224 = NISTObjectIdentifiers.id_sha224.getId();
    public static final String  DIGEST_SHA256 = NISTObjectIdentifiers.id_sha256.getId();
    public static final String  DIGEST_SHA384 = NISTObjectIdentifiers.id_sha384.getId();
    public static final String  DIGEST_SHA512 = NISTObjectIdentifiers.id_sha512.getId();
    public static final String  DIGEST_GOST3411 = CryptoProObjectIdentifiers.gostR3411.getId();
    public static final String  DIGEST_RIPEMD128 = TeleTrusTObjectIdentifiers.ripemd128.getId();
    public static final String  DIGEST_RIPEMD160 = TeleTrusTObjectIdentifiers.ripemd160.getId();
    public static final String  DIGEST_RIPEMD256 = TeleTrusTObjectIdentifiers.ripemd256.getId();	
	
    public final static Cryptographer Default = new Cryptographer();
    
    EncryptionAlgorithm m_encryptionAlgorithm;
    DigestAlgorithm m_digestAlgorithm;
    boolean m_includeEpilogue = true;

    public Cryptographer()
    {
        this.m_encryptionAlgorithm = EncryptionAlgorithm.RSA_3DES;
        this.m_digestAlgorithm = DigestAlgorithm.SHA1;
    }

    public Cryptographer(EncryptionAlgorithm encryptionAlgorithm, DigestAlgorithm digestAlgorithm)
    {
        this.m_encryptionAlgorithm = encryptionAlgorithm;
        this.m_digestAlgorithm = digestAlgorithm;
    }

    public EncryptionAlgorithm getEncryptionAlgorithm()
    {
        return this.m_encryptionAlgorithm;
    }
    
    public void setEncryptionAlgorithm(EncryptionAlgorithm value)
    {
        this.m_encryptionAlgorithm = value;
    }

    public DigestAlgorithm getDigestAlgorithm()
    {
    	return this.m_digestAlgorithm;
    }
    
    public void setDigestAlgorithm(DigestAlgorithm value)
    {
        this.m_digestAlgorithm = value;
    }
    
    /// <summary>
    /// When signing multipart messages, some mail clients do not include the multipart epilogue
    /// </summary>
    public boolean isIncludeMultipartEpilogueInSignature()
    {
        return this.m_includeEpilogue;
    }            
    
    public void setIncludeMultipartEpilogueInSignature(boolean value)
    {
        this.m_includeEpilogue = value;
    }
    
    //-----------------------------------------------------
    //
    // Encryption
    //
    //-----------------------------------------------------
    public MimeEntity encrypt(MultipartEntity entity, X509Certificate encryptingCertificate)
    {
        return this.encrypt(entity.toEntity(), encryptingCertificate);
    }

    public MimeEntity encrypt(MultipartEntity entity, Collection<X509Certificate> encryptingCertificates)
    {
        return this.encrypt(entity.toEntity(), encryptingCertificates);
    }
    
    public MimeEntity encrypt(MimeEntity entity, X509Certificate encryptingCertificate)
    {
    	Collection<X509Certificate> certs = new ArrayList<X509Certificate>();
    	certs.add(encryptingCertificate);
    	
        return this.encrypt(entity, certs);
    }

    public MimeEntity encrypt(MimeEntity entity,  Collection<X509Certificate> encryptingCertificates)
    {
        if (entity == null)
        {
            throw new IllegalArgumentException();
        }

        byte[] messageBytes = EntitySerializer.Default.serializeToBytes(entity);     // Serialize message out as ASCII encoded...
        	
        MimeBodyPart partToEncrypt = null;
        
        try
        {        	
        	partToEncrypt = new MimeBodyPart(new ByteArrayInputStream(messageBytes));
        }
        catch (Exception e)
        {
        	e.printStackTrace();
        }
        
        MimeBodyPart encryptedPart =  this.encrypt(partToEncrypt, encryptingCertificates);

        messageBytes = null;
        
        MimeEntity encryptedEntity = new MimeEntity();
        encryptedEntity.setContentType(new Header(Protocol.ContentTypeHeader, Protocol.EncryptedContentTypeHeaderValue));
        encryptedEntity.setContentTransferEncoding(new Header(Protocol.ContentTransferEncodingHeader, Protocol.TransferEncodingBase64));
                    

        
        try
        {
            InputStream stream = new BufferedInputStream(encryptedPart.getInputStream());
            
    	    ByteArrayOutputStream bos = new ByteArrayOutputStream();
    	    byte[] buffer = new byte[2048];
    	    int len;
    	    while ((len = stream.read(buffer, 0, 2048)) != -1)
    	    	bos.write(buffer, 0, len);
    	    
    	    bos.flush();
    	    
    	    buffer = bos.toByteArray();
        	
        	
        	encryptedEntity.setBody(new Body(Base64.encodeBase64String(buffer)));
        }
        catch (Exception e)
        {        
        	throw new ProtocolException(ProtocolError.Unexpected, e);
        }

        return encryptedEntity;
    }
    
    public MimeBodyPart encrypt(MimeBodyPart bodyPart, X509Certificate encryptingCertificate)
    {
    	Collection<X509Certificate> certs = new ArrayList<X509Certificate>();
    	certs.add(encryptingCertificate);    	
    	
        return this.encrypt(bodyPart, certs);
    }

    public MimeBodyPart encrypt(MimeBodyPart bodyPart, Collection<X509Certificate> encryptingCertificates)
    {
        return this.createEncryptedEnvelope(bodyPart, encryptingCertificates);
    }
            
    public MimeBodyPart createEncryptedEnvelope(MimeBodyPart bodyPart, Collection<X509Certificate> encryptingCertificates)
    {
        if (bodyPart == null || encryptingCertificates == null || encryptingCertificates.size() == 0)
        {
            throw new IllegalArgumentException();
        }

        SMIMEEnvelopedGenerator gen = new SMIMEEnvelopedGenerator();

        for(X509Certificate cert : encryptingCertificates)
        	gen.addKeyTransRecipient(cert);
        
        MimeBodyPart retVal = null;
        
        try
        {
        	retVal =  gen.generate(bodyPart, toEncyAlgorithmOid(this.m_encryptionAlgorithm), "BC");
        }
        catch (Exception e)
        {
        	throw new ProtocolException(ProtocolError.Unexpected, e);
        }
        
        return retVal;
    }

    //-----------------------------------------------------
    //
    // Decryption
    //
    //-----------------------------------------------------
    //
    // Cryptography is performed only on the Mime portions of the message, not the RFC822 headers
    //
    public MimeEntity decrypt(Message message, X509CertificateEx decryptingCertificate)
    {
        return this.decrypt(message.extractMimeEntity(), decryptingCertificate);
    }
    
    public MimeEntity decrypt(MimeEntity encryptedEntity, X509CertificateEx decryptingCertificate)
    {
        if (encryptedEntity == null || decryptingCertificate == null)
        {
            throw new IllegalArgumentException();
        }
        
        if (!decryptingCertificate.hasPrivateKey())
        {
            throw new IllegalArgumentException("Certificate has no private key");
        }
        
        encryptedEntity.verifyContentType(Protocol.EncryptedContentTypeHeaderValue);
        encryptedEntity.verifyTransferEncoding(Protocol.TransferEncodingBase64);
        
    	Collection<X509CertificateEx> certs = new ArrayList<X509CertificateEx>();
    		certs.add(decryptingCertificate);
    		
        MimeEntity retVal = this.decrypt(encryptedEntity, certs);
        
        //
        // And turn the decrypted bytes back into an entity
        //
        return retVal;
    }
    
    public MimeEntity decrypt(MimeEntity encryptedEntity/*byte[] encryptedContent*/, Collection<X509CertificateEx> decryptingCertificates)
    {
        if (decryptingCertificates == null || decryptingCertificates.size() == 0)
        {
            throw new IllegalArgumentException();
        }

        // just use the first for now
        byte[] decryptedPayload = null;
        MimeEntity retEntity = null;
        try
        {
        	byte[] encryptedContent = encryptedEntity.getBody().getText().getBytes();//Base64.decodeBase64(encryptedEntity.getBody().getText());
        	
        		        
        	InternetHeaders headers = new InternetHeaders();
        	for (Header header : encryptedEntity.getHeaders())
        		headers.addHeader(header.getName(), header.getValue());
        	
        	
        	MimeBodyPart msg = new MimeBodyPart(headers,encryptedContent);
            SMIMEEnveloped m = new SMIMEEnveloped(msg);            
            
            X509CertificateEx decryptCert = decryptingCertificates.iterator().next();
            
            RecipientId recId = new RecipientId();        	
	        recId.setSerialNumber(decryptCert.getSerialNumber());
	        recId.setIssuer(decryptCert.getIssuerX500Principal().getEncoded());
	
	        RecipientInformationStore recipients = m.getRecipientInfos();
	        RecipientInformation recipient = recipients.get(recId);	
	        	        	        	       

	        decryptedPayload = recipient.getContent(decryptCert.getPrivateKey(), "BC");
	        retEntity = EntitySerializer.Default.deserialize(decryptedPayload);
	        
        }
        catch (Exception e)
        {
        	e.printStackTrace();
        	throw new ProtocolException(ProtocolError.Unexpected);
        }
	    

        return retEntity;
    }


    //-----------------------------------------------------
    //
    // Signatures
    //
    //-----------------------------------------------------
    //
    // Cryptography is performed only on the Mime portions of the message, not the RFC822 headers
    // Some mail readers ignore the epilogue when calculating signatures!
    //
    public SignedEntity sign(Message message, X509CertificateEx signingCertificate)
    {
        return this.sign(message.extractEntityForSignature(this.m_includeEpilogue), signingCertificate);
    }
    
    public SignedEntity sign(MimeEntity entity, X509CertificateEx signingCertificate)        
    {
    	Collection<X509CertificateEx> certs = new ArrayList<X509CertificateEx>();
    	certs.add(signingCertificate);
    	
        return this.sign(entity, certs);
    }
    
    public SignedEntity sign(MimeEntity entity, Collection<X509CertificateEx> signingCertificates)
    {
        if (entity == null)
        {
            throw new IllegalArgumentException();
        }

        byte[] messageBytes = EntitySerializer.Default.serializeToBytes(entity);     // Serialize message out as ASCII encoded...
     
        MimeMultipart mm = this.createSignatureEntity(messageBytes, signingCertificates);
        SignedEntity retVal = null;
        
        try
        {
        
        	retVal = new SignedEntity(new ContentType(mm.getContentType()), mm);
        }
        catch (ParseException e)
        {
        	throw new ProtocolException(ProtocolError.InvalidHeader, e);
        }
        
        return retVal;
    }

    public MimeMultipart createSignatureEntity(byte[] entity, Collection<X509CertificateEx> signingCertificates)
    {    	
    	MimeMultipart retVal = null;
    	try
    	{
	        MimeBodyPart signedContent = new MimeBodyPart(new ByteArrayInputStream(entity));
	    		        
	    	ASN1EncodableVector signedAttrs = new ASN1EncodableVector();
	    	SMIMECapabilityVector caps = new SMIMECapabilityVector();
	
	    	caps.addCapability(SMIMECapability.dES_EDE3_CBC);
	    	caps.addCapability(SMIMECapability.rC2_CBC, 128);
	    	caps.addCapability(SMIMECapability.dES_CBC);
	    	caps.addCapability(new DERObjectIdentifier("1.2.840.113549.1.7.1"));
	    	caps.addCapability(PKCSObjectIdentifiers.x509Certificate);
	    	signedAttrs.add(new SMIMECapabilitiesAttribute(caps));  
	    	
	    	List  certList = new ArrayList();
	    	CMSSignedDataGenerator generator = new CMSSignedDataGenerator();
	    	for (X509CertificateEx signer : signingCertificates)
	    	{
	        	generator.addSigner(signer.getPrivateKey(), signer,
	          	      CMSSignedDataGenerator.DIGEST_SHA1, new AttributeTable(signedAttrs), null);
	        	certList.add(signer);
	    	}    	  	    		    	
	    	
	    	CertStore certsAndcrls = CertStore.getInstance("Collection", new CollectionCertStoreParameters(certList), "BC");   
	    	generator.addCertificatesAndCRLs(certsAndcrls);
	    	CMSProcessableBodyPartInbound content = new CMSProcessableBodyPartInbound(signedContent);
	    	
	    	CMSSignedData signedData = generator.generate(content, true, "BC");
	    	  	    	
	        String  header = new String("signed; protocol=\"application/pkcs7-signature\"; micalg=sha1");           
	        
	        String encodedSig = Base64.encodeBase64String(signedData.getEncoded());
	        
	        retVal = new MimeMultipart(header.toString());
	        
	        MimeBodyPart sig = new MimeBodyPart(new InternetHeaders(), encodedSig.getBytes("ASCII"));
            sig.addHeader("Content-Type", "application/pkcs7-signature; name=smime.p7s; smime-type=signed-data");
            sig.addHeader("Content-Disposition", "attachment; filename=\"smime.p7s\"");
            sig.addHeader("Content-Description", "S/MIME Cryptographic Signature");
            sig.addHeader("Content-Transfer-Encoding", "base64");
	        
            retVal.addBodyPart(signedContent);
            retVal.addBodyPart(sig);

    	}   
    	catch (Exception e)
    	{
    		throw new NHINDException(ProtocolError.Unexpected, e);    		
    	}    	
    	return retVal;  
  	
    }
    


    //-----------------------------------------------------
    //
    // Signature Validation
    //
    //-----------------------------------------------------
    
    public void checkSignature(SignedEntity signedEntity, X509Certificate signerCertificate, Collection<X509Certificate> anchors) throws SignatureValidationException
    {
    	CMSSignedData signatureEnvelope = deserializeSignatureEnvelope(signedEntity);
    	    	
    	try
    	{
	    	for (SignerInformation sigInfo : (Collection<SignerInformation>)signatureEnvelope.getSignerInfos().getSigners())
	    	{	    		
	    		sigInfo.verify(signerCertificate, "BC");
	    	}
    	}
    	catch (Exception e)
    	{
    		throw new SignatureValidationException("Signature validation failure.");
    	}
    }                        
 	
    public CMSSignedData deserializeSignatureEnvelope(SignedEntity entity)
    {

    	
    	if (entity == null)
        {
            throw new NHINDException();
        }

    	CMSSignedData signed = null;
    	
    	try
    	{
    		//signed = new SMIMESigned(entity.getMimeMultipart());
    		byte[] messageBytes = EntitySerializer.Default.serializeToBytes(entity.getContent());
            MimeBodyPart signedContent = null;
            
           	signedContent = new MimeBodyPart(new ByteArrayInputStream(messageBytes));

            signed = new CMSSignedData(new CMSProcessableBodyPartInbound(signedContent), entity.getMimeMultipart().getBodyPart(1).getInputStream());                        
    	}
    	catch (Exception e)
    	{
    		e.printStackTrace();
        	throw new ProtocolException(ProtocolError.Unexpected, e);
    	}
    	
    	return signed;
    }

    //
    // OIDs
    //
    static final String Oid_SHA1 = new String("1.3.14.3.2.26");
    static final String Oid_SHA256 = new String("2.16.840.1.101.3.4.2.1");
    static final String Oid_SHA512 = new String("2.16.840.1.101.3.4.2.3");
    static final String Oid_RSA_ThreeDES = new String("1.2.840.113549.3.7");
    static final String Oid_ContentType_Data = new String("1.2.840.113549.1.7.1");



    String toDigestAlgorithmOid(DigestAlgorithm type)
    {
        switch (type)
        {
            default:
                throw new IllegalArgumentException();

            case SHA1:
                return Cryptographer.Oid_SHA1;
        }
    }

    String toEncyAlgorithmOid(EncryptionAlgorithm type)
    {
        switch (type)
        {
            default:
                throw new IllegalArgumentException();

            case RSA_3DES:
                return SMIMEEnvelopedGenerator.DES_EDE3_CBC;
        }
    }
    
}
